package com.fulu.game.bigdata.realtime.sink



import com.alibaba.fastjson.JSONObject
import com.fulu.game.bigdata.realtime.config.Config
import com.fulu.game.bigdata.realtime.entity.{Field, Sink}
import com.google.common.base.Strings
import org.apache.flink.api.common.functions.RuntimeContext
import org.apache.flink.streaming.connectors.elasticsearch.{ElasticsearchSinkFunction, RequestIndexer}
import org.apache.flink.streaming.connectors.elasticsearch6.{ElasticsearchSink, RestClientFactory}
import org.apache.flink.types.Row
import org.apache.http.HttpHost
import org.apache.http.auth.{AuthScope, UsernamePasswordCredentials}
import org.apache.http.client.config.RequestConfig
import org.apache.http.impl.client.BasicCredentialsProvider
import org.apache.http.impl.nio.client.HttpAsyncClientBuilder
import org.elasticsearch.action.delete.DeleteRequest
import org.elasticsearch.action.index.IndexRequest
import org.elasticsearch.action.update.UpdateRequest
import org.elasticsearch.client.RestClientBuilder
import org.elasticsearch.common.xcontent.XContentType

import scala.collection.immutable.HashMap
import scala.collection.JavaConverters._


object ElasticsearchSinkBuilder {

  def createElasticSearchSink(sink: Sink): ElasticsearchSink[(Boolean, Row)] = {

    val esName = sink.getEsName

    val config: HashMap[String, String] = new HashMap[String,String]
    config.+ ("cluster.name" -> "elasticsearch")

    val hostAddress = Config.getElasticsearchHost(esName).split(",").toStream
      .map(hostAndPort => hostAndPort.split(":"))
      .map(item => new HttpHost(item(0),Integer.parseInt(item(1)))).toList.asJava

    val builder = new ElasticsearchSink.Builder[(Boolean, Row)](hostAddress,new ElasticsearchSinkFunction[(Boolean, Row)] {
      override def process(element: (Boolean, Row), ctx: RuntimeContext, indexer: RequestIndexer): Unit = {
        if (element._1) {
          val row = element._2
          val document = new JSONObject()
          for (i <- 0 until sink.getFields.size()) {
            val value = row.getField(i)
            document.put(sink.getFields.get(i).name, value)
          }
          val primaryKey = sink.getFields.stream().filter(field => field.isPrimaryKey).findFirst()
          if (primaryKey.isPresent) {
            val keyIndex = sink.getFields.indexOf(primaryKey.get())
            val key = String.valueOf(row.getField(keyIndex))
            val updateRequest = new UpdateRequest(sink.getEsIndex, sink.getEsType, key)
              .doc(document.toJSONString.getBytes(), XContentType.JSON)
              .upsert(document.toJSONString.getBytes(), XContentType.JSON)
            indexer.add(updateRequest)
          } else {
            val indexRequest = new IndexRequest(sink.getEsIndex, sink.getEsType).source(document.toJSONString, XContentType.JSON)
            indexer.add(indexRequest)
          }
        } else {
          val row = element._2
          val key = String.valueOf(row.getField(0))
          val deleteRequest = new DeleteRequest(sink.getEsIndex, sink.getEsType, key)
          indexer.add(deleteRequest)
        }
      }
    })

    if(!Strings.isNullOrEmpty(Config.getElasticsearchUsername(esName))) {
      builder.setRestClientFactory(new RestClientFactory {
        override def configureRestClientBuilder(restClientBuilder: RestClientBuilder): Unit = {
          restClientBuilder.setHttpClientConfigCallback(new RestClientBuilder.HttpClientConfigCallback {
            override def customizeHttpClient(httpClientBuilder: HttpAsyncClientBuilder): HttpAsyncClientBuilder = {
              val credentialsProvider = new BasicCredentialsProvider()
              credentialsProvider.setCredentials(AuthScope.ANY, new UsernamePasswordCredentials(Config.getElasticsearchUsername(esName), Config.getElasticsearchPassword(esName)))
              httpClientBuilder.setDefaultCredentialsProvider(credentialsProvider)
            }

          })
          restClientBuilder.setRequestConfigCallback(requestConfigBuilder => requestConfigBuilder.setConnectionRequestTimeout(300000).setSocketTimeout(600000).setConnectionRequestTimeout(0))
        }
      })
    }

    builder.build()
  }
}
